import glob, shutil, os, re, logging
from toolbox import update_ui, trimmed_format_exc, gen_time_str
from toolbox import CatchException, report_exception, get_log_folder
from toolbox import write_history_to_file, promote_file_to_downloadzone
import time
fast_debug = False

# 用户可在此设置最大token限制，用于切割长文本
MAX_TOKEN_LIMIT = 1024
# 用户可在此设置翻译工作的并发数量
# 双卡14 ，四卡30-40左右
MAX_WORKERS = 30
# 提示词
# 当前已经更新到v0.5版提示词，：进一步本土化翻译效果。
PROMPT = """        
一、翻译要求：
翻译MD源代码（包含latex源码）的内容翻译为目标语言。默认目标语言为中文。
中文语法，句式要求生动活泼，通俗易懂，句子流畅度要最高。句式温暖亲切，清新自然。
一切没有文采的，不生动的语句逐行润色成生动、活泼的句子。
一切死板、冰冷、复杂晦涩的句式逐行转化为通俗易懂的语句。
一切中文里没有的句式和表达转化为中文本土句式并保证通俗易懂。
仔细思考每一句话在中文里是如何恰当表达的，不要机械的翻译，输出的译文要完全像一个中文本土作者的作品。
只逐句翻译，不要有任何一个字的多余回答，比如解释和说明语句以及英文原文。
段落中的核心要点、重点概念的词语给与加粗显示。不要加粗整个段落，整个句子。

二、排版，格式要求：
appendix也要翻译。
去掉英文原文：输出翻译时不要带任何依据的英文原文。
输出的纯净性：保证输出的文本只包含纯净的译文，不输出其他任何译文之外的文本。
翻译完整性：保证每一句英文原文都要被翻译。
布局：译文要保持原文的框架结构和布局，保持原文语句先后和段落的先后顺序。
保留标题：译文保持和原文章标题的格式一一对应。
重点加黑：重点内容必须是短句或者词语才能加黑加粗，不能是长句子或者段落。

三、风险代码规避
保留数学公式：保留所有的数学公式的内容，不要删除。
latex规格化：原文中latex数学公式标签统一转换为`$` 和 `$$`这标准的标签样式，（只针对标签修改，和匹配，务必不要更改、删除公式内容）.
latex公式纠正：latex公式有渲染错误和语法错误也必须要纠正为正规的，通用的，能正常渲染的公式。
非代码的文本务必不要写到代码块里，保持普通正文的格式。
非标题的文本务必不要写到标题标签中，保持普通正文的格式。
图片丢失：审查对比原文章中的图片链接，保证译文中图片链接一个也不要丢失。

接下来请翻译：

"""

class PaperFileGroup():
    def __init__(self):
        self.file_paths = []
        self.file_contents = []
        self.sp_file_contents = []
        self.sp_file_index = []
        self.sp_file_tag = []

        # count_token
        from request_llms.bridge_all import model_info
        enc = model_info["gpt-3.5-turbo"]['tokenizer']
        def get_token_num(txt): return len(enc.encode(txt, disallowed_special=()))
        self.get_token_num = get_token_num

    def run_file_split(self, max_token_limit=MAX_TOKEN_LIMIT):
        """
        将长文本分离开来
        """
        for index, file_content in enumerate(self.file_contents):
            if self.get_token_num(file_content) < max_token_limit:
                self.sp_file_contents.append(file_content)
                self.sp_file_index.append(index)
                self.sp_file_tag.append(self.file_paths[index])
            else:
                from crazy_functions.pdf_fns.breakdown_txt import breakdown_text_to_satisfy_token_limit
                segments = breakdown_text_to_satisfy_token_limit(file_content, max_token_limit)
                for j, segment in enumerate(segments):
                    self.sp_file_contents.append(segment)
                    self.sp_file_index.append(index)
                    self.sp_file_tag.append(self.file_paths[index] + f".part-{j}.md")
        logging.info('Segmentation: done')

    def merge_result(self):
        self.file_result = ["" for _ in range(len(self.file_paths))]
        for r, k in zip(self.sp_file_result, self.sp_file_index):
            self.file_result[k] += r

    def write_result(self, language):
        manifest = []
        for path, res in zip(self.file_paths, self.file_result):
            import re
            base_dir = os.path.dirname(path)
            folder_name = os.path.basename(base_dir)
            base_name = 'full-中文翻译'
            index = 1
            dst_file = os.path.join(base_dir, f'{base_name}.md')
            while os.path.exists(dst_file):
                dst_file = os.path.join(base_dir, f'{base_name}-{index}.md')
                index += 1
            with open(dst_file, 'w', encoding='utf8') as f:
                manifest.append(dst_file)
                f.write(res)
            valid_folder_name = re.sub(r'[\\/*?"<>|]', '_', folder_name)
            if index > 1:
                folder_file = os.path.join(base_dir, f'{valid_folder_name}-{index}.md')
            else:
                folder_file = os.path.join(base_dir, f'{valid_folder_name}.md')
            with open(folder_file, 'w', encoding='utf8') as f:
                manifest.append(folder_file)
                f.write(res)
        return manifest

# 在多文件翻译函数中添加统计和报告功能
def 多文件翻译(file_manifest, project_folder, llm_kwargs, plugin_kwargs, chatbot, history, system_prompt, language='en'):
    from .crazy_utils import request_gpt_model_multi_threads_with_very_awesome_ui_and_high_efficiency

    #  <-------- 读取Markdown文件，删除其中的所有注释 ---------->
    pfg = PaperFileGroup()

    for index, fp in enumerate(file_manifest):
        # 输出统计信息
        total_count = len(file_manifest)
        current_index = index + 1
        folder_name = os.path.basename(os.path.dirname(fp))
        chatbot.append((f"待翻译总数: {total_count}, 当前序号: {current_index}, 当前文件夹名: {folder_name}", "开始翻译"))
        yield from update_ui(chatbot=chatbot, history=history) # 刷新界面
        pfg = PaperFileGroup()
        with open(fp, 'r', encoding='utf-8', errors='replace') as f:
            file_content = f.read()
            # 记录删除注释后的文本
            pfg.file_paths.append(fp)
            pfg.file_contents.append(file_content)

        #  <-------- 拆分过长的Markdown文件 ---------->
        pfg.run_file_split(max_token_limit=MAX_TOKEN_LIMIT)
        n_split = len(pfg.sp_file_contents)

        #  <-------- 多线程翻译开始 ---------->
        if language == 'en->zh':
            # 构建发送给大语言模型的输入内容列表
            # 每个输入内容由提示词和待翻译的文本片段组成，中间用两个换行符分隔
            inputs_array = [
                PROMPT + f"\n\n{frag}" for frag in pfg.sp_file_contents
            ]
            inputs_show_user_array = [f"翻译 {f}" for f in pfg.sp_file_tag]
            sys_prompt_array = ["You are a professional academic paper translator." + plugin_kwargs.get("additional_prompt", "") for _ in range(n_split)]
        elif language == 'zh->en':
            # 构建发送给大语言模型的输入内容列表
            # 每个输入内容由提示词和待翻译的文本片段组成，中间用两个换行符分隔
            inputs_array = [
                PROMPT + f"\n\n{frag}" for frag in pfg.sp_file_contents
            ]
            inputs_show_user_array = [f"翻译 {f}" for f in pfg.sp_file_tag]
            sys_prompt_array = ["You are a professional academic paper translator." + plugin_kwargs.get("additional_prompt", "") for _ in range(n_split)]
        else:
            # 构建发送给大语言模型的输入内容列表
            # 每个输入内容由提示词和待翻译的文本片段组成，中间用两个换行符分隔
            inputs_array = [
                PROMPT + f"\n\n{frag}" for frag in pfg.sp_file_contents
            ]
            inputs_show_user_array = [f"翻译 {f}" for f in pfg.sp_file_tag]
            sys_prompt_array = ["You are a professional academic paper translator." + plugin_kwargs.get("additional_prompt", "") for _ in range(n_split)]

        gpt_response_collection = yield from request_gpt_model_multi_threads_with_very_awesome_ui_and_high_efficiency(
            inputs_array=inputs_array,
            inputs_show_user_array=inputs_show_user_array,
            llm_kwargs=llm_kwargs,
            chatbot=chatbot,
            history_array=[[""] for _ in range(n_split)],
            sys_prompt_array=sys_prompt_array,
            max_workers=MAX_WORKERS,  # OpenAI所允许的最大并行过载
            scroller_max_len = 80
        )
        try:
            pfg.sp_file_result = []
            for i_say, gpt_say in zip(gpt_response_collection[0::2], gpt_response_collection[1::2]):
                if '</think>' in gpt_say:
                    gpt_say = gpt_say.split('</think>')[-1]
                pfg.sp_file_result.append(gpt_say)
            pfg.merge_result()
            output_file_arr = pfg.write_result(language)
            for output_file in output_file_arr:
                promote_file_to_downloadzone(output_file, chatbot=chatbot)
                if 'markdown_expected_output_path' in plugin_kwargs:
                    expected_f_name = plugin_kwargs['markdown_expected_output_path']
                    shutil.copyfile(output_file, expected_f_name)
            chatbot.append((f"{fp} 翻译成功", "翻译完成"))
            yield from update_ui(chatbot=chatbot, history=history) # 刷新界面
        except:
            logging.error(trimmed_format_exc())
    #  <-------- 整理结果，退出 ---------->
    create_report_file_name = gen_time_str() + f"-chatgpt.md"
    res = write_history_to_file(gpt_response_collection, file_basename=create_report_file_name)
    promote_file_to_downloadzone(res, chatbot=chatbot)
    history = gpt_response_collection
    chatbot.append((f"所有文件翻译完成", res))
    yield from update_ui(chatbot=chatbot, history=history) # 刷新界面


    # 统计和报告功能
    subdirectories = [d for d in os.listdir(project_folder) if os.path.isdir(os.path.join(project_folder, d))]
    translated_subdirectories = []
    untranslated_subdirectories = []
    problematic_subdirectories = []
    for subdir in subdirectories:
        subdir_path = os.path.join(project_folder, subdir)
        # 这里简单假设包含翻译结果文件的目录为已翻译目录，可根据实际情况修改判断逻辑
        if any(f.endswith('-中文翻译.md') for f in os.listdir(subdir_path)):
            translated_subdirectories.append(subdir_path)
        else:
            untranslated_subdirectories.append(subdir_path)
    # 可以根据实际情况添加对有问题目录的判断逻辑
    # 生成报告文件
    report_dir = os.path.join(os.getcwd(), 'reports')
    if not os.path.exists(report_dir):
        os.makedirs(report_dir)
    timestamp = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
    report_file = os.path.join(report_dir, f'{timestamp}-translation_report.md')
    with open(report_file, 'w', encoding='utf-8') as f:
        f.write('---\n')
        f.write(f'# 数字统计\n')
        f.write('---\n')
        f.write(f'深度为1层的子目录个数: {len(subdirectories)}\n')
        f.write(f'翻译的子目录个数: {len(translated_subdirectories)}\n')
        f.write(f'没翻译的子目录个数: {len(untranslated_subdirectories)}\n')
        f.write('---\n')
        f.write(f'# 翻译结果明细\n')
        f.write('---\n')
        f.write('## 被翻译过的目录:\n')
        for dir_path in translated_subdirectories:
            f.write(f'- {dir_path}\n')
        f.write('## 没有被翻译的目录:\n')
        for dir_path in untranslated_subdirectories:
            f.write(f'- {dir_path}\n')
        f.write('## 其他有问题的目录:\n')
        for dir_path in problematic_subdirectories:
            f.write(f'- {dir_path}\n')
    promote_file_to_downloadzone(report_file, chatbot=chatbot)


def get_files_from_everything(txt, preference=''):
    if txt == "": return False, None, None
    success = True
    if txt.startswith('http'):
        import requests
        from toolbox import get_conf
        proxies = get_conf('proxies')
        # 网络的远程文件
        if preference == 'Github':
            logging.info('正在从github下载资源 ...')
            if not txt.endswith('.md'):
                # Make a request to the GitHub API to retrieve the repository information
                url = txt.replace("https://github.com/", "https://api.github.com/repos/") + '/readme'
                response = requests.get(url, proxies=proxies)
                txt = response.json()['download_url']
            else:
                txt = txt.replace("https://github.com/", "https://raw.githubusercontent.com/")
                txt = txt.replace("/blob/", "/")

        r = requests.get(txt, proxies=proxies)
        download_local = f'{get_log_folder(plugin_name="批量Markdown翻译")}/raw-readme-{gen_time_str()}.md'
        project_folder = f'{get_log_folder(plugin_name="批量Markdown翻译")}'
        with open(download_local, 'wb+') as f: f.write(r.content)
        file_manifest = [download_local]
    elif txt.endswith('.md'):
        # 直接给定文件
        file_manifest = [txt]
        project_folder = os.path.dirname(txt)
    elif os.path.exists(txt):
        # 本地路径，递归搜索
        project_folder = txt
        file_manifest = [f for f in glob.glob(f'{project_folder}/**/full.md', recursive=True)]
        if len(file_manifest) == 0:
            success = False
    else:
        project_folder = None
        file_manifest = []
        success = False

    return success, file_manifest, project_folder


@CatchException
def Markdown英译中(txt, llm_kwargs, plugin_kwargs, chatbot, history, system_prompt, user_request):
    # 基本信息：功能、贡献者
    chatbot.append([
        "函数插件功能？",
        "对整个Markdown项目进行翻译。函数插件贡献者: Binary-Husky"])
    yield from update_ui(chatbot=chatbot, history=history) # 刷新界面

    # 尝试导入依赖，如果缺少依赖，则给出安装建议
    try:
        import tiktoken
    except:
        report_exception(chatbot, history,
                         a=f"解析项目: {txt}",
                         b=f"导入软件依赖失败。使用该模块需要额外依赖，安装方法```pip install --upgrade tiktoken```。")
        yield from update_ui(chatbot=chatbot, history=history) # 刷新界面
        return
    history = []    # 清空历史，以免输入溢出

    success, file_manifest, project_folder = get_files_from_everything(txt, preference="Github")

    if not success:
        # 什么都没有
        if txt == "": txt = '空空如也的输入栏'
        report_exception(chatbot, history, a = f"解析项目: {txt}", b = f"找不到本地项目或无权访问: {txt}")
        yield from update_ui(chatbot=chatbot, history=history) # 刷新界面
        return

    if len(file_manifest) == 0:
        report_exception(chatbot, history, a = f"解析项目: {txt}", b = f"找不到任何.md文件: {txt}")
        yield from update_ui(chatbot=chatbot, history=history) # 刷新界面
        return

    yield from 多文件翻译(file_manifest, project_folder, llm_kwargs, plugin_kwargs, chatbot, history, system_prompt, language='en->zh')





@CatchException
def Markdown中译英(txt, llm_kwargs, plugin_kwargs, chatbot, history, system_prompt, user_request):
    # 基本信息：功能、贡献者
    chatbot.append([
        "函数插件功能？",
        "对整个Markdown项目进行翻译。函数插件贡献者: Binary-Husky"])
    yield from update_ui(chatbot=chatbot, history=history) # 刷新界面

    # 尝试导入依赖，如果缺少依赖，则给出安装建议
    try:
        import tiktoken
    except:
        report_exception(chatbot, history,
                         a=f"解析项目: {txt}",
                         b=f"导入软件依赖失败。使用该模块需要额外依赖，安装方法```pip install --upgrade tiktoken```。")
        yield from update_ui(chatbot=chatbot, history=history) # 刷新界面
        return
    history = []    # 清空历史，以免输入溢出
    success, file_manifest, project_folder = get_files_from_everything(txt)
    if not success:
        # 什么都没有
        if txt == "": txt = '空空如也的输入栏'
        report_exception(chatbot, history, a = f"解析项目: {txt}", b = f"找不到本地项目或无权访问: {txt}")
        yield from update_ui(chatbot=chatbot, history=history) # 刷新界面
        return
    if len(file_manifest) == 0:
        report_exception(chatbot, history, a = f"解析项目: {txt}", b = f"找不到任何.md文件: {txt}")
        yield from update_ui(chatbot=chatbot, history=history) # 刷新界面
        return
    yield from 多文件翻译(file_manifest, project_folder, llm_kwargs, plugin_kwargs, chatbot, history, system_prompt, language='zh->en')


@CatchException
def Markdown翻译指定语言(txt, llm_kwargs, plugin_kwargs, chatbot, history, system_prompt, user_request):
    # 基本信息：功能、贡献者
    chatbot.append([
        "函数插件功能？",
        "对整个Markdown项目进行翻译。函数插件贡献者: Binary-Husky"])
    yield from update_ui(chatbot=chatbot, history=history) # 刷新界面

    # 尝试导入依赖，如果缺少依赖，则给出安装建议
    try:
        import tiktoken
    except:
        report_exception(chatbot, history,
                         a=f"解析项目: {txt}",
                         b=f"导入软件依赖失败。使用该模块需要额外依赖，安装方法```pip install --upgrade tiktoken```。")
        yield from update_ui(chatbot=chatbot, history=history) # 刷新界面
        return
    history = []    # 清空历史，以免输入溢出
    success, file_manifest, project_folder = get_files_from_everything(txt)
    if not success:
        # 什么都没有
        if txt == "": txt = '空空如也的输入栏'
        report_exception(chatbot, history, a = f"解析项目: {txt}", b = f"找不到本地项目或无权访问: {txt}")
        yield from update_ui(chatbot=chatbot, history=history) # 刷新界面
        return
    if len(file_manifest) == 0:
        report_exception(chatbot, history, a = f"解析项目: {txt}", b = f"找不到任何.md文件: {txt}")
        yield from update_ui(chatbot=chatbot, history=history) # 刷新界面
        return

    if ("advanced_arg" in plugin_kwargs) and (plugin_kwargs["advanced_arg"] == ""): plugin_kwargs.pop("advanced_arg")
    language = plugin_kwargs.get("advanced_arg", 'Chinese')
    yield from 多文件翻译(file_manifest, project_folder, llm_kwargs, plugin_kwargs, chatbot, history, system_prompt, language=language)